import numpy as np
from PIL import Image
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Activation, Flatten, Dense
from keras.optimizers import SGD
from keras.src.utils import np_utils
from sklearn import datasets
from sklearn.model_selection import train_test_split

# Config values
num_classes = 10
img_depth = 1
img_height = 28
img_width = 28

def build_lenet(input_shape, num_classes):
    """
    Build the LeNet model.
    :param input_shape: (height, width, depth)
    :param num_classes: the classes number
    :return: the built model
    """
    # Creating the LeNet model
    model = Sequential()

    # Adding the first convolutional layer
    model.add(Conv2D(20, (5, 5), padding="same", input_shape=input_shape))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    # Adding the second convolutional layer
    model.add(Conv2D(50, (5, 5), padding="same"))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

    # Adding the fully connected layers
    model.add(Flatten())
    model.add(Dense(500))
    model.add(Activation("relu"))

    # Adding a softmax layer
    model.add(Dense(num_classes))
    model.add(Activation("softmax"))

    return model

def predict_image(model, image_path):
    """
    Use the trained model to predict image.
    :param model: the trained LeNet model
    :param image_path: path to your image
    :return: the predict result
    """
    # load image
    img = Image.open(image_path).convert("L")  # convert to grayscale
    img = img.resize((img_width, img_height))
    img_array = np.array(img, dtype="float32") / 255.0 # normalization
    img_array = img_array.reshape(1, img_height, img_width, img_depth)

    # 预测类别
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions)
    confidence = predictions[0][predicted_class]
    return predicted_class, confidence

# Load MNIST data
mnist = datasets.fetch_openml("mnist_784", version=1)
data = mnist.data.to_numpy().reshape((mnist.data.shape[0], img_height, img_width, img_depth))
data = data.astype("float32") / 255.0  # Normalize the images to [0, 1.0]

# Convert labels to integers
labels = mnist.target.astype("int")

# Split the data into train and test set
train_data, test_data, train_label, test_label = train_test_split(
    data, labels, test_size=0.25, random_state=42
)
train_label = np_utils.to_categorical(train_label, num_classes)
test_label = np_utils.to_categorical(test_label, num_classes)

# build LeNet model
input_shape = (img_height, img_width, img_depth)
model = build_lenet(input_shape, num_classes)

# Set the loss functions and evaluation metrics
model.compile(loss="categorical_crossentropy", optimizer=SGD(learning_rate=0.01), metrics=["accuracy"])

# Train the LeNet model
model.fit(train_data, train_label, batch_size=32, epochs=30, verbose=1)

# Test the model
loss, accuracy = model.evaluate(test_data, test_label, batch_size=64, verbose=1)
print("Accuracy: {:.2f}%".format(accuracy * 100))

# predict your image
image_path = "images/8.png"
predicted_class, confidence = predict_image(model, image_path)
print(f"Predicted Class: {predicted_class}, Confidence: {confidence:.2f}")